import { FlappyBirdEnv } from './environment.js';
import { PPOAgent } from './ppo_agent.js';

let env;
let agent;
let experiences = [];
let state;
let episodeReward = 0;
let lastActionProb = 0.5; // 存储最后一次动作概率

// 速度控制
let speedSlider;
let speedValue;

async function initTF() {
    await tf.ready();
    console.log('TensorFlow.js initialized');
}

async function init() {
    await initTF();
    
    const canvas = createCanvas(800, 600);
    canvas.parent('gameCanvas');
    frameRate(60);
    
    // 初始化速度控制
    speedSlider = document.getElementById('speedSlider');
    speedValue = document.getElementById('speedValue');
    speedSlider.addEventListener('input', function() {
        speedValue.textContent = `${this.value}x`;
    });
    
    env = new FlappyBirdEnv();
    env.width = width;
    env.height = height;
    
    agent = new PPOAgent(4, 2);
    state = env.reset();
}

function setup() {
    init().catch(console.error);
}

async function draw() {
    if (!env || !agent) return;
    
    try {
        const iterations = parseInt(speedSlider.value);
        
        for (let i = 0; i < iterations; i++) {
            const action = await agent.getAction(state);
            const [nextState, reward, done] = env.step(action);
            
            experiences.push({
                state: state,
                action: action,
                reward: reward,
                nextState: nextState,
                done: done
            });

            episodeReward += reward;
            state = nextState;

            if (experiences.length >= agent.batchSize) {
                const trainingBatch = experiences.slice(0, agent.batchSize);
                await agent.train(trainingBatch);
                experiences = experiences.slice(agent.batchSize);
            }

            if (done) {
                console.log(`Episode finished with reward: ${episodeReward}`);
                state = env.reset();
                episodeReward = 0;
                break;
            }
        }

        env.render();
        
    } catch (error) {
        console.error('Error in draw loop:', error);
    }
}

window.setup = setup;
window.draw = draw; 